import matplotlib.pyplot as plt
import networkx as nx
import torch
from scipy.sparse import lil_matrix


def create_met_png(x_axis, y_axis, data):
    # fig = plt.figure(dpi=600)
    fig = plt.figure()
    # 画图
    plt.plot(x_axis, y_axis[0], label='before attack')
    plt.plot(x_axis, y_axis[1], label='after attack')

    # 设置标题、x轴标签、y轴标签、图例位置
    plt.title('mettack on ' + data)
    plt.xlabel('disturbing rate')
    plt.ylabel('accuracy')
    plt.legend(loc='best')
    return fig.canvas

    # 显示图像
    # plt.show()


def create_png(x_axis, y_axis, data, data_type, algorithm):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5.5))

    # 绘制子图1
    ax1.plot(x_axis, y_axis[0], label='before attack')
    ax1.plot(x_axis, y_axis[1], label='features attack')
    ax1.plot(x_axis, y_axis[3], label='mix attack')
    title1 = data + ' with ' + data_type + " on " + algorithm
    ax1.set_title(title1)
    ax1.set_xlabel('disturbing rate')
    ax1.set_ylabel('accuracy')
    ax1.legend(loc='best')

    # 绘制子图2
    ax2.plot(x_axis, y_axis[0], label='before attack')
    ax2.plot(x_axis, y_axis[2], label='structure attack')
    ax2.plot(x_axis, y_axis[3], label='mix attack')
    title2 = data + ' with ' + data_type + " on " + algorithm
    ax2.set_title(title2)
    ax2.set_xlabel('disturbing rate')
    ax2.set_ylabel('accuracy')
    ax2.legend(loc='best')

    return fig.canvas


def whole_picture(algorithm, dataset, data_type):
    result = {
        "mettack":
            {"cora":
                 [[0.8346, 0.8342, 0.8136, 0.7626, 0.5445, 0.3147],
                  [0.8436, 0.8426, 0.8441, 0.8415, 0.8441, 0.8446]],
             "citeseer":
                 [[0.7365, 0.7306, 0.7148, 0.6658, 0.4963, 0.2659],
                  [0.7536, 0.7494, 0.753, 0.7487, 0.753, 0.7494]],
             "x_axis": ['2%', '5%', '10%', '30%', '50%', '80%']
             },
        "nettack":
            {
                "cora":
                    {
                        "train":
                            {
                                "features": [0.8275, 0.8265, 0.8199, 0.8270, 0.8305, 0.8270, 0.8315, 0.8179, 0.8270,
                                             0.8205],
                                "structure": [0.8199, 0.8043, 0.7998, 0.7853, 0.7792, 0.7757, 0.7656, 0.7545, 0.7490,
                                              0.7409],
                                "mix": [0.8265, 0.8154, 0.8058, 0.7938, 0.7857, 0.7626, 0.7384, 0.7359, 0.7269, 0.6972]
                            },
                        "test":
                            {
                                "features": [0.8280, 0.8147, 0.7988, 0.7897, 0.7767, 0.7666, 0.7772, 0.7681, 0.7736,
                                             0.7671],
                                "structure": [0.7490, 0.6786, 0.6192, 0.5951, 0.5493, 0.5226, 0.4935, 0.4251, 0.4195,
                                              0.4256],
                                "mix": [0.7344, 0.6509, 0.5991, 0.5573, 0.5161, 0.5010, 0.4608, 0.4180, 0.3793, 0.3657]
                            },
                        "y_axis": [0.8405, 0.8399, 0.8414, 0.8395, 0.8436, 0.8426, 0.8441, 0.8425, 0.8441, 0.8436]
                    },
                "citeseer":
                    {
                        "train":
                            {
                                "features": [0.7382, 0.7352, 0.7340, 0.7322, 0.7299, 0.7269, 0.7251, 0.7127, 0.7269,
                                             0.7097],
                                "structure": [0.7316, 0.7239, 0.7204, 0.7156, 0.6973, 0.6866, 0.6795, 0.6736, 0.6646,
                                              0.6564],
                                "mix": [0.6938, 0.6813, 0.6700, 0.6469, 0.5936, 0.5729, 0.5641, 0.5569, 0.5551, 0.5415]
                            },
                        "test":
                            {
                                "features": [0.7429, 0.7376, 0.6831, 0.6629, 0.6605, 0.6682, 0.6434, 0.6339, 0.6143,
                                             0.6013],
                                "structure": [0.6570, 0.6007, 0.5557, 0.5160, 0.4479, 0.4017, 0.3827, 0.3632, 0.3483,
                                              0.3477],
                                "mix": [0.6434, 0.5983, 0.5450, 0.5036, 0.4277, 0.3857, 0.3620, 0.3406, 0.3116, 0.2956]
                            },
                        "y_axis": [0.7536, 0.7494, 0.753, 0.7487, 0.753, 0.7494, 0.7518, 0.7482, 0.7500, 0.7493]
                    }
            },
        "ig_attack":
            {
                "cora":
                    {
                        "train":
                            {
                                "features": [0.8395, 0.8350, 0.8355, 0.8345, 0.8350, 0.8355, 0.8375, 0.8385, 0.8365,
                                             0.8375],
                                "structure": [0.8305, 0.8234, 0.7912, 0.7867, 0.7807, 0.7731, 0.7601, 0.7560, 0.7389,
                                              0.7128],
                                "mix": [0.8300, 0.8169, 0.8058, 0.7953, 0.7872, 0.7741, 0.7606, 0.7505, 0.7354, 0.7022]
                            },
                        "test":
                            {
                                "features": [0.8365, 0.8280, 0.8295, 0.8229, 0.8255, 0.8204, 0.8144, 0.8099, 0.8027,
                                             0.8073],
                                "structure": [0.8038, 0.7792, 0.7515, 0.7193, 0.7018, 0.6840, 0.6643, 0.6598, 0.6549,
                                              0.6403],
                                "mix": [0.7895, 0.7702, 0.7367, 0.7148, 0.6862, 0.6629, 0.6451, 0.6437, 0.6178, 0.5941]
                            },
                        "y_axis": [0.8405, 0.8399, 0.8414, 0.8395, 0.8436, 0.8426, 0.8441, 0.8425, 0.8441, 0.8436]
                    },
                "citeseer":
                    {
                        "train":
                            {
                                "features": [0.7441, 0.7387, 0.7346, 0.7387, 0.7347, 0.7305, 0.7328, 0.7328, 0.7316,
                                             0.7340],
                                "structure": [0.7393, 0.7287, 0.7227, 0.7150, 0.6789, 0.6700, 0.6653, 0.6558, 0.6374,
                                              0.6351],
                                "mix": [0.7263, 0.7020, 0.6914, 0.6878, 0.6789, 0.6605, 0.6540, 0.6357, 0.6333, 0.6327]
                            },
                        "test":
                            {
                                "features": [0.7383, 0.7325, 0.7335, 0.7251, 0.7235, 0.7241, 0.7177, 0.7204, 0.7156,
                                             0.7135],
                                "structure": [0.6825, 0.6503, 0.6019, 0.5793, 0.5698, 0.5487, 0.5366, 0.5197, 0.5161,
                                              0.4934],
                                "mix": [0.6898, 0.6340, 0.6103, 0.5903, 0.5534, 0.5429, 0.5319, 0.5134, 0.5266, 0.5013]
                            },
                        "y_axis": [0.7536, 0.7494, 0.753, 0.7487, 0.753, 0.7494, 0.7518, 0.7482, 0.75, 0.7493]
                    }
            },
        "x_axis": ['10%', '20%', '30%', '40%', '50%', '60%', '70%', '80%', '90%', '100%']
    }
    if algorithm == "mettack":
        return create_met_png(x_axis=result[algorithm]['x_axis'], y_axis=result[algorithm][dataset], data=dataset)
    else:
        return create_png(x_axis=result['x_axis'], y_axis=[result[algorithm][dataset]['y_axis'],
                                                           result[algorithm][dataset][data_type]["features"],
                                                           result[algorithm][dataset][data_type]["structure"],
                                                           result[algorithm][dataset][data_type]["mix"]],
                          data=dataset, data_type=data_type, algorithm=algorithm)


def fix_lil_matrix(lil_mat):
    """
    Takes a lil_matrix as input and returns a new lil_matrix
    with fixed rows and data lists (same length as rows).
    """
    fixed_mat = lil_matrix(lil_mat.shape, dtype=lil_mat.dtype)
    for i, row in enumerate(lil_mat.rows):
        fixed_mat.rows[i] = row
        num_cols = len(row)
        if num_cols > len(lil_mat.data[i]):
            fixed_mat.data[i] = lil_mat.data[i] + [1.0] * (num_cols - len(lil_mat.data[i]))
        else:
            fixed_mat.data[i] = lil_mat.data[i][:num_cols]
    return fixed_mat


def single_pictue(algorithm, dataset, node):
    """
    可视化节点，默认向下跳1个节点
    """
    adj = torch.load('./save_data/' + dataset + '/' + dataset + '_adj.pt')
    modified_adj = torch.load('./outputs/' + algorithm + '/' + dataset + '/test/structure/100%_nodes_modified_adj.pt')
    # 使用邻接矩阵创建一个 Graph 对象
    G = nx.from_scipy_sparse_matrix(adj)
    # 找个节点，对他进行可视化
    # 获取节点 source 的所有一跳邻居节点
    neighbors1 = nx.single_source_shortest_path_length(G, node, cutoff=1).keys()
    subgraph1 = G.subgraph(neighbors1)

    modified_adj = fix_lil_matrix(modified_adj)
    modified_G = nx.from_scipy_sparse_matrix(modified_adj)
    neighbors2 = nx.single_source_shortest_path_length(modified_G, node, cutoff=1).keys()
    subgraph2 = modified_G.subgraph(neighbors2)

    color_map = {0: 'orange', 1: 'green', 2: 'red'}
    # 默认情况下，节点展示为橙色
    node_colors1 = {i: color_map[0] for i in subgraph1.nodes}
    node_colors2 = {i: color_map[0] for i in subgraph2.nodes}
    if len(subgraph1.nodes) > len(subgraph2.nodes):
        reduce_node = subgraph1.nodes - subgraph2.nodes
        for i in reduce_node:
            # 减少的节点，用红色表示
            node_colors1[i] = color_map[2]
    else:
        add_node = subgraph2.nodes - subgraph1.nodes
        for i in add_node:
            # 增加的节点，用绿色表示
            node_colors2[i] = color_map[1]
    # 把字典里的颜色取出来
    node_colors1 = [node_colors1[i] for i in node_colors1]
    node_colors2 = [node_colors2[i] for i in node_colors2]

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5.5))
    pos1 = nx.shell_layout(subgraph1)
    pos2 = nx.shell_layout(subgraph2)

    # 设置节点大小，指定的节点，就显示大一点
    node_size = 900
    selected_node_size = 2000
    node_sizes1 = [selected_node_size if i == node else node_size for i in subgraph1.nodes]
    node_sizes2 = [selected_node_size if i == node else node_size for i in subgraph2.nodes]

    # 给子图加上标题
    ax1.set_title("Before node %d is attacked by %s" % (node, algorithm))
    ax2.set_title("After Node %d is attacked by %s" % (node, algorithm))
    # 绘制图形
    nx.draw_networkx(subgraph1, pos=pos1, node_size=node_sizes1, node_color=node_colors1, with_labels=True, ax=ax1)
    nx.draw_networkx(subgraph2, pos=pos2, node_size=node_sizes2, node_color=node_colors2, with_labels=True, ax=ax2)
    # plt.show()
    return fig.canvas

# single_pictue('ig_attack', 'cora', 490)
